# 导入必要的库
import os
import numpy as np
import logging
from collections import Counter
from sklearn.model_selection import train_test_split
from lazypredict.Supervised import LazyClassifier
from util import createXY
from train_ensemble import (
    voting_ensemble,
    bagging_ensemble,
    stacking_ensemble,
    pasting_ensemble,
    adaboost_ensemble,
    gradient_boosting_ensemble,
    evaluate_individual_classifiers,
    save_models,
    display_results_table
)

# 配置logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 配置参数
def get_config():
    """
    获取配置参数
    """
    config = {
        'feature': 'flat',  # 固定使用flat特征提取
        'train_dir': '../data/train',  # 训练集目录
        'dest_dir': '.',  # 输出目录
        'use_lazypredict': True,  # 是否启用Lazypredict
        'save_results': True  # 是否保存结果
    }
    return config

# 使用Lazypredict评估模型
def evaluate_with_lazypredict(X_train, X_test, y_train, y_test, save_results=True, dest_dir="."):
    """
    使用Lazypredict快速评估多种模型的性能，并保存结果到文件（可选）。
    
    参数:
    X_train, X_test, y_train, y_test: 数据集
    save_results: 是否保存评估结果
    dest_dir: 保存结果的目录
    """
    logging.info("开始使用Lazypredict评估模型...")
    clf = LazyClassifier(verbose=0, ignore_warnings=True, custom_metric=None)
    models, predictions = clf.fit(X_train, X_test, y_train, y_test)
    logging.info("Lazypredict模型评估完成。")
    
    # 打印模型评估结果
    print(models)
    
    # 保存结果到文件
    if save_results:
        results_path = os.path.join(dest_dir, "lazypredict_results.csv")
        models.to_csv(results_path)
        logging.info(f"Lazypredict评估结果已保存到 {results_path}")
    
    return models

# 主函数
def main():
    config = get_config()
    
    logging.info(f"特征提取方法: {config['feature'].upper()}")
    logging.info(f"训练集目录: {config['train_dir']}")
    logging.info(f"缓存输出目录: {config['dest_dir']}")

    # 载入和预处理数据
    X, y = createXY(train_folder=config['train_dir'], dest_folder=config['dest_dir'], method=config['feature'])
    X = np.array(X).astype('float32')
    y = np.array(y)
    logging.info("数据加载和预处理完成。")
    logging.info(f"数据集大小: {X.shape}, 标签分布: {Counter(y)}")

    # 数据集分割
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=2023)
    logging.info("数据集划分为训练集和测试集。")
    logging.info(f"训练集大小: {X_train.shape}, 测试集大小: {X_test.shape}")

    # 使用Lazypredict评估模型
    if config['use_lazypredict']:
        evaluate_with_lazypredict(X_train, X_test, y_train, y_test, save_results=config['save_results'], dest_dir=config['dest_dir'])

    # 运行集成学习方法
    all_results = {}
    voting_results = voting_ensemble(X_train, y_train, X_test, y_test)
    all_results['voting'] = voting_results
    
    bagging_results = bagging_ensemble(X_train, y_train, X_test, y_test)
    all_results['bagging'] = bagging_results
    
    stacking_results = stacking_ensemble(X_train, y_train, X_test, y_test)
    all_results['stacking'] = stacking_results
    
    pasting_results = pasting_ensemble(X_train, y_train, X_test, y_test)
    all_results['pasting'] = pasting_results

    adaboost_results = adaboost_ensemble(X_train, y_train, X_test, y_test)
    all_results['adaboost'] = adaboost_results

    gradient_boosting_results = gradient_boosting_ensemble(X_train, y_train, X_test, y_test)
    all_results['gradient_boosting'] = gradient_boosting_results

    # 显示结果表格
    display_results_table(all_results)
    
    # 保存模型
    if config['save_results']:
        save_models(all_results, config['dest_dir'])
        logging.info("所有模型已成功保存！")

if __name__ == '__main__':
    main()
